"""
训练主程序 - 启动自我对弈训练
"""

import torch
import argparse
from self_play_trainer import SelfPlayTrainer


def main():
    parser = argparse.ArgumentParser(description='五子棋自我对弈训练')
    parser.add_argument('--games', type=int, default=100, help='对弈局数')
    parser.add_argument('--simulations', type=int, default=100, help='MCTS模拟次数')
    parser.add_argument('--batch-size', type=int, default=32, help='训练批次大小')
    parser.add_argument('--save-interval', type=int, default=10, help='保存模型间隔')
    parser.add_argument('--model-path', type=str, default='models/best_model.pth', help='模型保存路径')
    parser.add_argument('--resume', type=str, default=None, help='继续训练的模型路径')
    
    args = parser.parse_args()
    
    # 检测设备
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"🖥️  使用设备: {device}")
    
    if device == 'cuda':
        print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
    
    # 创建训练器
    trainer = SelfPlayTrainer(model_path=args.resume, device=device)
    
    # 开始训练
    trainer.train(
        num_games=args.games,
        num_simulations=args.simulations,
        batch_size=args.batch_size,
        save_interval=args.save_interval,
        model_save_path=args.model_path
    )


if __name__ == "__main__":
    print("=" * 60)
    print("🎯 五子棋 AlphaGo Zero 训练系统")
    print("=" * 60)
    print()
    
    # 快速开始选项
    print("请选择训练模式:")
    print("1. 快速测试 (20局, 50次模拟) - 约5分钟")
    print("2. 标准训练 (100局, 100次模拟) - 约30分钟")
    print("3. 深度训练 (500局, 200次模拟) - 约3小时")
    print("4. 自定义参数")
    print()
    
    choice = input("请输入选择 (1-4): ").strip()
    
    if choice == '1':
        games, simulations = 20, 50
    elif choice == '2':
        games, simulations = 100, 100
    elif choice == '3':
        games, simulations = 500, 200
    elif choice == '4':
        games = int(input("对弈局数: "))
        simulations = int(input("MCTS模拟次数: "))
    else:
        print("❌ 无效选择，使用默认参数")
        games, simulations = 100, 100
    
    print()
    print("=" * 60)
    print(f"📊 训练配置:")
    print(f"  - 对弈局数: {games}")
    print(f"  - MCTS模拟次数: {simulations}")
    print(f"  - 批次大小: 32")
    print(f"  - 保存间隔: 每10局")
    print("=" * 60)
    print()
    
    # 检测设备
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"🖥️  使用设备: {device}")
    
    if device == 'cuda':
        print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
    
    print()
    input("按回车键开始训练...")
    print()
    
    # 创建训练器
    trainer = SelfPlayTrainer(device=device)
    
    # 开始训练
    trainer.train(
        num_games=games,
        num_simulations=simulations,
        batch_size=32,
        save_interval=10,
        model_save_path='models/best_model.pth'
    )
    
    print()
    print("=" * 60)
    print("✅ 训练完成！")
    print("=" * 60)
    print()
    print("下一步:")
    print("1. 运行 'python play_trained.py' 测试训练好的AI")
    print("2. 继续训练: 'python train.py' 并选择继续训练")
    print()
    input("按回车键退出...")
